import numpy as np
from collections import Counter

class SimpleKNN:
    def __init__(self, k=3):
        self.k = k  
        self.X_train = None
        self.y_train = None

    # 计算欧氏距离
    def _distance(self, x1, x2):
        return np.sqrt(np.sum((x1 - x2) ** 2, axis=1))

    # 训练
    def fit(self, X_train, y_train):
        self.X_train = X_train
        self.y_train = y_train

    # 预测
    def predict(self, X_test):
        predictions = []
        for x in X_test:
            dists = self._distance(x, self.X_train)

            k_neighbor_labels = self.y_train[np.argsort(dists)[:self.k]]

            predictions.append(Counter(k_neighbor_labels).most_common(1)[0][0])
        return np.array(predictions)

if __name__ == "__main__":

    X_train = np.array([[1,2], [2,3], [3,4], [4,5], [1,3], [4,3]])
    y_train = np.array([0, 0, 0, 1, 1, 1])
    
    # 测试样本
    X_test = np.array([[2,2], [3,3]])  
    
    # 运行KNN
    knn = SimpleKNN(k=3)
    knn.fit(X_train, y_train)
    y_pred = knn.predict(X_test)
    
    print("预测结果：", y_pred) 